package org.dromara.common.encrypt.filter;

import cn.hutool.core.util.RandomUtil;
import jakarta.servlet.ServletOutputStream;
import jakarta.servlet.WriteListener;
import jakarta.servlet.http.HttpServletResponse;
import jakarta.servlet.http.HttpServletResponseWrapper;
import org.dromara.common.encrypt.utils.EncryptUtils;

import java.io.*;
import java.nio.charset.StandardCharsets;


/**
 * @Author 刘武贵
 * @Date 2024/8/15 22:00
 * @Description 功能描述：加密响应参数包装类
 */
/**
 * 加密响应体包装类，继承自HttpServletResponseWrapper
 * 该类的目的是对服务器发送给客户端的响应体进行加密处理
 * 它通过包装实际的HttpServletResponse对象，拦截对响应体的写入操作，并将写入的数据进行加密后再实际写入
 * 这样做可以增加数据在传输过程中的安全性
 */
public class EncryptResponseBodyWrapper extends HttpServletResponseWrapper {

    /**
     * 将输出流转换为字节数组输出
     */
    private final ByteArrayOutputStream byteArrayOutputStream;

    /**
     * 用于向客户端输出数据
     */
    private final ServletOutputStream servletOutputStream;

    /**
     * 用于向客户端输出字符数据
     */
    private final PrintWriter printWriter;

    /**
     * 构造函数，用于初始化EncryptResponseBodyWrapper对象
     * 它接收一个HttpServletResponse对象，并对其进行包装，以便支持对响应体进行加密处理
     * 该构造函数主要负责初始化内部的输出流和打印写入器，以便后续能够捕获和处理响应数据
     *
     * @param response HttpServletResponse对象，代表HTTP响应
     * @throws IOException 如果在初始化过程中发生I/O错误
     */
    public EncryptResponseBodyWrapper(HttpServletResponse response) throws IOException {
        super(response);
        this.byteArrayOutputStream = new ByteArrayOutputStream();
        this.servletOutputStream = this.getOutputStream();
        this.printWriter = new PrintWriter(new OutputStreamWriter(byteArrayOutputStream));
    }

    /**
     * 获取与响应关联的PrintWriter对象
     * 通过此方法可以获得用于向客户端输出文本的流
     *
     * @return PrintWriter对象，用于向客户端输出文本
     */
    @Override
    public PrintWriter getWriter() {
        return printWriter;
    }

    /**
     * 刷新响应体的缓冲区
     * 此方法确保将缓冲区中的内容立即发送到客户端
     * 它首先检查是否存在servletOutputStream对象，如果存在，则调用其flush方法
     * 然后，如果存在printWriter对象，也调用其flush方法
     *
     * @throws IOException 如果在刷新过程中发生I/O错误
     */
    @Override
    public void flushBuffer() throws IOException {
        if (servletOutputStream != null) {
            servletOutputStream.flush();
        }
        if (printWriter != null) {
            printWriter.flush();
        }
    }

    /**
     * 重置输出流
     * 该方法清空当前的字节数组输出流，使其能够重新使用
     */
    @Override
    public void reset() {
        byteArrayOutputStream.reset();
    }

    /**
     * 获取HTTP响应数据
     *
     * @return 返回当前输出流中的字节数组
     * @throws IOException 如果在获取数据过程中发生I/O错误
     */
    public byte[] getResponseData() throws IOException {
        // 确保所有缓冲区的数据都被写出
        flushBuffer();
        // 将输出流转换为字节数组并返回
        return byteArrayOutputStream.toByteArray();
    }

    /**
     * 获取输出流内容作为字符串
     *
     * @return 返回输出流中的内容作为字符串
     * @throws IOException 如果在获取内容过程中发生I/O错误
     */
    public String getContent() throws IOException {
        // 确保所有缓冲区的数据都被写出
        flushBuffer();
        // 将输出流转换为字符串并返回
        return byteArrayOutputStream.toString();
    }

    /**
     * 获取加密内容
     * 使用混合加密技术加密内容
     *
     * @param servletResponse response HTTP响应对象，用于设置加密后的响应头和内容
     * @param publicKey       RSA公钥 (用于加密 AES 秘钥)
     * @param headerFlag      请求头标志 自定义响应头标志，用于标识加密后的AES密钥
     * @return 加密后的内容
     * @throws IOException 如果发生I/O错误
     */
    public String getEncryptContent(HttpServletResponse servletResponse, String publicKey, String headerFlag) throws IOException {
        // 生成秘钥 生成32位随机AES密钥
        String aesPassword = RandomUtil.randomString(32);
        // 秘钥使用 Base64 编码 使用Base64编码AES密钥，以便在HTTP响应中安全传输
        String encryptAes = EncryptUtils.encryptByBase64(aesPassword);
        // Rsa 公钥加密 Base64 编码 使用提供的RSA公钥加密Base64编码后的AES密
        String encryptPassword = EncryptUtils.encryptByRsa(encryptAes, publicKey);

        // 设置响应头，允许前端跨域请求访问加密后的内容
        servletResponse.addHeader("Access-Control-Expose-Headers", headerFlag);
        servletResponse.setHeader(headerFlag, encryptPassword);
        servletResponse.setHeader("Access-Control-Allow-Origin", "*");
        servletResponse.setHeader("Access-Control-Allow-Methods", "*");
        // 设置响应字符集为UTF-8
        servletResponse.setCharacterEncoding(StandardCharsets.UTF_8.toString());

        // 获取原始内容，准备加密
        String originalBody = this.getContent();
        // 对内容进行加密 使用生成的AES密钥加密原始内容
        return EncryptUtils.encryptByAes(originalBody, aesPassword);
    }

    /**
     * 获取与响应对象相关的输出流
     * 此方法返回一个实现了ServletOutputStream接口的匿名类，重写了其中的多个方法
     * 主要目的是将写入的操作委托给一个ByteArrayOutputStream实例，以便于收集输出数据
     *
     * @return ServletOutputStream 实现了ServletOutputStream接口的匿名类，用于处理输出流
     * @throws IOException 如果出现I/O错误
     */
    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        return new ServletOutputStream() {

            /**
             * 检查输出流是否准备好向客户端写入数据
             * 在这种情况下，始终返回false，表示流未准备好，这是因为实际的写入操作由ByteArrayOutputStream实例处理
             *
             * @return boolean 始终返回false，表示流未准备好
             */
            @Override
            public boolean isReady() {
                return false;
            }


            /**
             * 设置一个WriteListener，用于监听输出流的写操作
             * 在这种情况下，该方法为空实现，因为实际的写入操作由ByteArrayOutputStream实例处理
             *
             * @param writeListener 写入监听器，未使用
             */
            @Override
            public void setWriteListener(WriteListener writeListener) {

            }


            /**
             * 将指定的字节写入输出流
             * 实际操作是将字节写入ByteArrayOutputStream实例
             *
             * @param b 要写入的字节
             * @throws IOException 如果出现I/O错误
             */
            @Override
            public void write(int b) throws IOException {
                byteArrayOutputStream.write(b);
            }


            /**
             * 将整个字节数组写入输出流
             * 实际操作是将字节数组全部写入ByteArrayOutputStream实例
             *
             * @param b 要写入的字节数组
             * @throws IOException 如果出现I/O错误
             */
            @Override
            public void write(byte[] b) throws IOException {
                byteArrayOutputStream.write(b);
            }


            /**
             * 将字节数组的一部分写入输出流
             * 实际操作是将指定范围的字节写入ByteArrayOutputStream实例
             *
             * @param b 要写入的字节数组
             * @param off 起始偏移量
             * @param len 要写入的长度
             * @throws IOException 如果出现I/O错误
             */
            @Override
            public void write(byte[] b, int off, int len) throws IOException {
                byteArrayOutputStream.write(b, off, len);
            }
        };
    }

}
